import argparse
import os
import cv2
import numpy as np
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim
import lpips
import torch

def load_image(path, size=None):
    """加载图片并转为RGB格式（统一尺寸，避免计算误差）"""
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if len(img.shape) == 4:  # 处理RGBA图片（去除Alpha通道）
        img = img[..., :3]
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # BGR转RGB
    if size is not None:  # 强制统一尺寸
        img = cv2.resize(img, size, interpolation=cv2.INTER_LINEAR)
    return img.astype(np.float32) / 255.0  # 归一化到[0,1]

def calc_lpips(img1, img2, device):
    """计算LPIPS（感知指标）：输入为[0,1]的RGB图，输出越接近0越好"""
    # LPIPS要求输入格式：(1,3,H,W)，值域[-1,1]
    transform = lambda x: torch.from_numpy(x.transpose(2,0,1)).unsqueeze(0).to(device) * 2 - 1
    img1_tensor = transform(img1)
    img2_tensor = transform(img2)
    lpips_model = lpips.LPIPS(net='alex')  # alexnet为常用模型
    dist = lpips_model(img1_tensor, img2_tensor).item()
    return dist

def main():
    parser = argparse.ArgumentParser(description="计算HR与LR的PSNR/SSIM/LPIPS指标（批量处理）")
    parser.add_argument('--hr_dir', type=str, required=True, help="高分辨率图(HR)文件夹路径")
    parser.add_argument('--lr_dir', type=str, required=True, help="低分辨率图(LR)文件夹路径")
    parser.add_argument('--ext', type=str, default='auto', help="图片格式（auto/jpg/png，默认自动匹配）")
    args = parser.parse_args()

    # 设备选择（优先GPU，无则CPU）
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备：{device}")

    # 获取所有图片文件名（HR与LR文件名需一致）
    ext = args.ext if args.ext != 'auto' else ''
    hr_paths = sorted([p for p in os.listdir(args.hr_dir) if p.endswith(ext) or ext == ''])
    lr_paths = sorted([p for p in os.listdir(args.lr_dir) if p.endswith(ext) or ext == ''])

    # 校验HR与LR文件数量匹配
    if len(hr_paths) != len(lr_paths):
        raise ValueError(f"HR与LR文件数量不匹配！HR:{len(hr_paths)} | LR:{len(lr_paths)}")

    # 初始化结果列表
    results = []
    print(f"\n开始计算{len(hr_paths)}张图片的HR-LR指标...")
    print("="*60)
    print(f"{'图片名':<20} {'HR-LR PSNR':<12} {'HR-LR SSIM':<12} {'HR-LR LPIPS':<12}")
    print("="*60)

    for hr_name, lr_name in zip(hr_paths, lr_paths):
        # 拼接完整路径
        hr_path = os.path.join(args.hr_dir, hr_name)
        lr_path = os.path.join(args.lr_dir, lr_name)

        # 加载图片（计算时需统一尺寸，以HR为基准）
        hr_img = load_image(hr_path)
        lr_img = load_image(lr_path, size=(hr_img.shape[1], hr_img.shape[0]))  # 强制匹配HR尺寸

        # 计算 HR-LR 指标
        hr_lr_psnr = psnr(hr_img, lr_img, data_range=1.0)
        hr_lr_ssim = ssim(hr_img, lr_img, data_range=1.0, channel_axis=2)  # channel_axis=2（RGB图）
        hr_lr_lpips = calc_lpips(hr_img, lr_img, device)

        # 保存结果
        results.append({
            'img_name': hr_name,
            'hr_lr_psnr': hr_lr_psnr,
            'hr_lr_ssim': hr_lr_ssim,
            'hr_lr_lpips': hr_lr_lpips
        })

        # 打印单张图片结果
        print(f"{hr_name[:18]:<20} {hr_lr_psnr:<12.2f} {hr_lr_ssim:<12.4f} {hr_lr_lpips:<12.4f}")

    # 计算平均指标
    avg_hr_lr_psnr = np.mean([r['hr_lr_psnr'] for r in results])
    avg_hr_lr_ssim = np.mean([r['hr_lr_ssim'] for r in results])
    avg_hr_lr_lpips = np.mean([r['hr_lr_lpips'] for r in results])

    print("="*60)
    print(f"{'平均值':<20} {avg_hr_lr_psnr:<12.2f} {avg_hr_lr_ssim:<12.4f} {avg_hr_lr_lpips:<12.4f}")
    print("="*60)

    # 保存结果到文件
    output_file = os.path.join(args.lr_dir, 'hr_lr_metrics_summary.txt')  # 结果保存在LR文件夹
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("HR与LR指标计算结果\n")
        f.write(f"HR文件夹：{args.hr_dir}\n")
        f.write(f"LR文件夹：{args.lr_dir}\n")
        f.write("="*60 + "\n")
        f.write(f"{'图片名':<20} {'HR-LR PSNR':<12} {'HR-LR SSIM':<12} {'HR-LR LPIPS':<12}\n")
        f.write("="*60 + "\n")
        for r in results:
            f.write(f"{r['img_name'][:18]:<20} {r['hr_lr_psnr']:<12.2f} {r['hr_lr_ssim']:<12.4f} {r['hr_lr_lpips']:<12.4f}\n")
        f.write("="*60 + "\n")
        f.write(f"{'平均值':<20} {avg_hr_lr_psnr:<12.2f} {avg_hr_lr_ssim:<12.4f} {avg_hr_lr_lpips:<12.4f}\n")
    print(f"\n结果已保存到：{output_file}")

if __name__ == '__main__':
    main()